﻿Imports System.Runtime.CompilerServices
Imports System.Xml.Linq.Extensions

Namespace Extension
    Public Module TreeExtensions
        <Extension>
        Public Iterator Function ChildrenEx(item As DependencyObject) As IEnumerable(Of DependencyObject)
            Dim childrenCount = VisualTreeHelper.GetChildrenCount(item)
            For i = 0 To childrenCount - 1
                Yield VisualTreeHelper.GetChild(item, i)
            Next
        End Function

        <Extension>
        Public Function Ancestors(Of T)(item As DependencyObject) As IEnumerable(Of T)
            Return item.Ancestors.Where(Function(i) TypeOf i Is T).Cast(Of T)()
        End Function

        <Extension>
        Public Iterator Function Ancestors(item As DependencyObject) As IEnumerable(Of DependencyObject)
            Dim p = item.ParentEx
            While p IsNot Nothing
                Yield p
                p = p.ParentEx
            End While
        End Function

        <Extension>
        Public Function ParentEx(item As DependencyObject) As DependencyObject
            Return VisualTreeHelper.GetParent(item)
        End Function

        <Extension>
        Public Function Descendants(Of T)(item As DependencyObject) As IEnumerable(Of T)
            Return item.Descendants.Where(Function(i) TypeOf i Is T).Cast(Of T)()
        End Function

        <Extension>
        Public Iterator Function Descendants(item As DependencyObject) As IEnumerable(Of DependencyObject)
            For Each c In item.ChildrenEx
                Yield c
                For Each g In c.Descendants
                    Yield g
                Next
            Next
        End Function
    End Module
End Namespace